package cz.cuni.amis.episodic.lisp;

import cz.cuni.amis.episodic.lisp.behan.LispTree;
import cz.cuni.amis.episodic.lisp.behan.EventProcessor;
import cz.cuni.amis.episodic.lisp.behan.TraceToEventVisitor;
import cz.cuni.amis.episodic.lisp.netcreators.CHMMCreator;
import cz.cuni.amis.episodic.lisp.netcreators.NetCreator;
import cz.cuni.amis.episodic.lisp.netcreators.TraceModificationStrategy;
import cz.cuni.amis.episodic.lisp.visitor.TreeTraceVisitor;
import java.io.FileNotFoundException;
import java.io.FileOutputStream;
import java.io.OutputStream;
import java.io.PrintStream;
import java.util.*;

import jsmile.utils.Learn;
import lombok.Cleanup;
import org.apache.commons.lang.StringUtils;
import smile.Network;

/**
 * Converts Lisp plan to CSV format plus creates XDSL Dynamic Bayes Network.
 * 
 * @author ik
 */
public class LispPlanToCsvData {

	public static abstract class TreeToTraceTranslator implements EventProcessor {

		int traceDepthLimit;

		
		public TreeToTraceTranslator(int traceDepthLimit) {
			this.traceDepthLimit = traceDepthLimit;
		}

		List<String[]> traces = new ArrayList<>();

		public List<String[]> getTraces() {
			return traces;
		}

		@Override
		public boolean openNode(String goal, ArrayList<String> trace,
				ArrayList<String> affordances) {
			return true;
		}

		@Override
		public boolean closeNode(String goal, ArrayList<String> trace) {
			return true;
		}
		
		/**
		 * Hook that can be called after all traces were processed.
		 */
		public void postprocess() {
			
		}
		
	}

	public static class BehanFilteredProcessor extends TreeToTraceTranslator {

		public BehanFilteredProcessor(int traceDepthLimit) {
			super(traceDepthLimit);
		}

		@Override
		public boolean atomicAction(String action, ArrayList<String> trace,
				ArrayList<String> affordances) {

			int processedDepth;
			if (traceDepthLimit <= 0) {
				processedDepth = trace.size() / 2 + 1;
			} else {
				processedDepth = Math
						.min(trace.size() / 2 + 1, traceDepthLimit);
			}

			// filter AND nodes

			String[] filteredTrace = new String[processedDepth];
			int offset = trace.size() - (processedDepth - 1) * 2;
			for (int i = 0; i < processedDepth - 1; i++) {
				filteredTrace[i] = DefaultLispTranslator.translateOutcome(trace
						.get(offset + i * 2));
			}
			// store
			filteredTrace[filteredTrace.length - 1] = DefaultLispTranslator
					.translateOutcome(action);
			traces.add(filteredTrace);
			return true;
		}
	}
		
	public static class LispToCsvResult {

		public List<List<String>> levelActions = new ArrayList<>();
		public int lineNum = 0;
		public int batchSize = -1;
		public String inputLispPlan;
	}

	public static LispToCsvResult lispToCsv(String inputLispPlan,
			OutputStream outputPlanCsv, int batchSize,
			TreeToTraceTranslator translator, TraceModificationStrategy strategy) throws FileNotFoundException {
		LispTree parser = new LispTree(inputLispPlan);
		LispToCsvResult res = new LispToCsvResult();

		test(parser);
		parser.walkInForest(new TraceToEventVisitor(null, null, null,
				translator));

		// TODO strategy should be passed to translator ...
		if(strategy != null) {
			translator.traces = strategy.modifyTrace(translator.traces, true);
		}
		
		// overwrite default batch size with the new computed from data 
		if(strategy != null && strategy.getBatchSize() > 0) batchSize = strategy.getBatchSize();
		
		// find max depth
		int max = Integer.MIN_VALUE;
		for (String[] trace : translator.getTraces()) {
			max = Math.max(max, trace.length);
		}

		@Cleanup
		PrintStream out = new PrintStream(outputPlanCsv);
		String delim = "\t";

		// each set contains actions found on the corresponding level in
		// hierarchy
		Set<String>[] levels = new HashSet[max];
		for (int i = 0; i < max; i++) {
			levels[i] = new HashSet<String>();
		}

		// print header
		String[] nodesInHeader = new String[max];
		for (int i = 0; i < nodesInHeader.length - 1; i++) {
			nodesInHeader[i] = "G" + (max - i - 1);
		}
		nodesInHeader[nodesInHeader.length - 1] = "O";
		List<String> header = new ArrayList<>(Arrays.asList("Batch", "Id"));
		header.addAll(Arrays.asList(nodesInHeader));
		out.println(StringUtils.join(header, delim));

		// output with padding

		int batchNum = 1;
		int linesInBatch = 0;
		String[] line = new String[max];
		for (String[] trace : translator.getTraces()) {
			linesInBatch++;
			out.print(batchNum + delim + Integer.toString(res.lineNum++)
					+ delim);
			int padCount = max - trace.length;
			for (int i = 0; i < padCount; i++) {
				out.print(trace[0]);
				if (i < max - 1) {
					out.print(delim);
				}
				levels[i].add(trace[0]);
			}
			out.println(StringUtils.join(trace, delim));

			if (batchSize > 0 && linesInBatch == batchSize) {
				batchNum++;
				linesInBatch = 0;
			}

			// remember actions/goals found in each level
			for (int i = 0; i < trace.length; i++) {
				levels[i + padCount].add(trace[i]);
			}
		}

		// create result
		res.batchSize = batchSize;
		res.inputLispPlan = inputLispPlan;

		// sort actions by alphabet
		for (int i = 0; i < levels.length; i++) {
			List<String> actions = new ArrayList<>(levels[i]);
			Collections.sort(actions);
			res.levelActions.add(actions);
		}
		return res;
	}

	public static void test(LispTree parser) {
		parser.walkInForest(new TreeTraceVisitor() {

			@Override
			public void traceVisited(List<List<String>> trace) {
				List<String> episTrace = new ArrayList<>();
				for (List<String> elem : trace) {
					episTrace.add(elem.get(0));
				}
				System.out.println(StringUtils.join(episTrace, " -> "));
				
			}
			
		});
	}

	/*
	 * public static void lispToDataAndNet(String inputLispPlan, String
	 * outputPlanCsv, String outputNetworkFilename, int maxDepthLimit,
	 * NetArchitecture architecture) throws FileNotFoundException { }
	 * 
	 * public static void lispToDataAndNet(String inputLispPlan, OutputStream
	 * outputPlanCsv) throws FileNotFoundException { }
	 */
	public static Network createDBNFromCsvResult(LispToCsvResult result,
			NetCreator netCreator, String outputNetworkFilename,
			String smileDataFile) {
		Network net = netCreator.createNet(result.levelActions);
		net.setDescription("Goal recognition network created from LISP plan file "
				+ result.inputLispPlan);
		net.setName("Goal recognition for " + result.inputLispPlan);

		// net.setSliceCount(result.batchSize > 0 ? result.batchSize :
		// result.lineNum);

/*		int last = outputNetworkFilename.lastIndexOf(".");
		String[] split = outputNetworkFilename.split("\\.");
		net.writeFile(outputNetworkFilename.substring(0, last) + "_before.xdsl");
		*/
		// learn the network
		if (smileDataFile != null)
			Learn.learnEM(net, smileDataFile);

		// write the net
		net.writeFile(outputNetworkFilename);
                return net;
	}

	public static void main(String[] args) throws FileNotFoundException {
		final int hiddenNodeStatesNum = 3;
		final String outputNetworkFilename = "behan_d3.xdsl";
		final String inputLispPlan = "src/main/resources/episplanner/testdata1.lisp";
		final String outputPlanCsv = "testdata1-d3.csv";
		final NetCreator netCreator = new CHMMCreator(false);
		final int maxDepthLimit = 3;
		// size of cases used for learning and at the same time number of
		// temporal slices
		// if it is <= 0 then maximum is used
		final int batchSize = 15;

		LispToCsvResult res = lispToCsv(inputLispPlan, new FileOutputStream(
				outputPlanCsv), batchSize, new BehanFilteredProcessor(
				maxDepthLimit), null);

		// create network
		createDBNFromCsvResult(res, netCreator, outputNetworkFilename, null);
	}
}
